from langchain_openai import ChatOpenAI
import settings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.tools import tool
from datetime import datetime

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_redis import RedisChatMessageHistory
from langchain_core.runnables import ConfigurableFieldSpec

from langchain_core.messages import AIMessage, ToolMessage

from .memory import ZLRedisChatMessageHistory
from .tools import ZLToolManager


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    return ZLRedisChatMessageHistory(session_id=session_id,
                                     redis_url=f"redis://{settings.REDIS_HOST}:{settings.REDIS_PORT}")


class ZLAgent:
    def __init__(self):
        self.tool_manager = ZLToolManager()
        self.tool_manager.star_scheduler()
        self.tool_manager.connect_signal(self.initialize)

    async def initialize(self, *args, **kwargs):
        self.llm = ChatOpenAI(
            base_url=settings.LLM_BASE_URL,
            model=settings.LLM_MODEL_NAME,
            api_key=settings.LLM_API_KEY,
            streaming=True
        )
        system_prompt = """
            你的名字叫小婵同学，是基于钉钉生态环境的个人助理，专心回答用户问题，不要杜撰不存在的信息
             * 工具调用：
                    * 当用户要求的操作或问题中涉及基于当前时间时，请先通过current_time工具获取当前时间。
                    * 涉及到将时间转换为ISO-8601格式时，要在末尾加上+08:00，代表的是东八区的北京时间。
                * 同义词替代：
                    * 当提示词中出现：任务、待办事项、工作内容、todo，都属于待办。

        """

        prompt_template = ChatPromptTemplate([
            ('system', system_prompt),
            ('placeholder', "{chat_history}"),
            ('user', "{input}"),
            ('placeholder', "{agent_scratchpad}")
        ])

        agent = create_tool_calling_agent(
            llm=self.llm,
            tools=await self.tool_manager.get_tools(),
            prompt=prompt_template
        )

        self.agent_executor = AgentExecutor(
            max_iterations=100,
            agent=agent,
            tools=await self.tool_manager.get_tools(),
            verbose=True,
            return_intermediate_steps=True
        )

    async def run(self, question, sender_staff_id: str):
        assert self.agent_executor is not None
        question += f"sender_staff_id: {sender_staff_id}"
        history = get_session_history(sender_staff_id)
        history.add_user_message(question)
        response = await self.agent_executor.ainvoke(
            {"input": question, "chat_history": history.messages},
            config={"configurable": {"session_id": sender_staff_id}}
        )

        steps = response.get('intermediate_steps', [])
        for step in steps:
            tool_call_action, tool_call_result = step
            for message in tool_call_action.message_log:
                tool_calls = message.tool_calls
                for tool_call in tool_calls:
                    if tool_call['id'] == tool_call_action.tool_call_id:
                        # 2.1. 添加大模型选择了哪个工具的消息
                        ai_message = AIMessage(content="", tool_calls=[tool_call])
                        history.add_message(ai_message)
                        tool_message = ToolMessage(content=str(tool_call_result),
                                                   tool_call_id=tool_call_action.tool_call_id)
                        history.add_message(tool_message)
                        break

        history.add_ai_message(response['output'])
        return response['output']
